Apply function to points within circular neighborhood by ahijevyc · Pull Request #941 · UXARRAY/uxarray

ahijevyc · 2024-09-09T16:44:16Z

Apply a neighborhood filter within a circular radius r to a UxDataset or UxDataArray.

Closes #930

Overview

This is kind of like uxarray.UxDataArray.inverse_distance_weighted_remap , but the neighborhood is defined by distance, not a number of nearest neighbors. This is ideally suited for a variable resolution mesh, in which a constant of neighbors doesn't have a constant sized neighborhood. Another difference is that this neighborhood filter does not weight data by inverse distance.

Just like uxarray.UxDataArray.subset.bounding_circle this function uses ball_tree.query_radius to select grid elements in a circular neighborhood, but this function finds the neighborhood for all elements in grid, not just one center_coordinate.

The filter function func may be a user-defined function, but uses np.mean by default. It could be min, max, np.median. It can even use functions that require additional arguments, like np.percentile if you supply the argument(s) with functools.partial (see below)

Expected Usage

from functools import partial
import numpy as np
import uxarray

grid_path = "/glade/campaign/mmm/wmr/weiwang/cps/irma3/2020/tk707_conus/init.nc"
data_path = "/glade/campaign/mmm/wmr/weiwang/cps/irma3/mp6/tk707/diag.2017-09-07_09.00.00.nc"
uxds = uxarray.open_mfdataset(
    grid_path,
    data_path
)

# Trim domain
lon_bounds = (-74, -64)
lat_bounds = (18, 24)
uxda = uxds["refl10cm_max"].isel(Time=0).subset.bounding_box(lon_bounds, lat_bounds)

# this is how you use this function to smooth with 0.25-deg filter.
uxda_mean = uxda.neighborhood_filter(func=np.mean, r=0.25)


# this is another way to use this function with np.percentile
uxda_max = uxda.neighborhood_filter(func=partial(np.percentile, q=90), r=0.25)

(uxda.plot.rasterize() + uxda_mean.plot.rasterize() + uxda_max.plot.rasterize()).cols(1)

PR Checklist

General

An issue is linked created and linked
Add appropriate labels
Filled out Overview and Expected Usage (if applicable) sections

Testing

Adequate tests are created if there is new functionality
Tests cover all possible logical paths in your function
Tests are not too basic (such as simply calling a function and nothing else)

Documentation

Docstrings have been added to all new functions
Docstrings have updated with any function changes
Internal functions have a preceding underscore (_); _neighborhood_filter is internal to uxarray/grid/neighbors.py
User functions added to docs/api.rst (the split user/internal api files no longer exist)

Examples

Any new notebook examples added to docs/examples/ folder
Clear the output of all cells before committing
New notebook files added to docs/examples.rst toctree
New notebook files added to new entry in docs/gallery.yml with appropriate thumbnail photo in docs/_static/thumbnails/

philipc2

A few initial comments:

Kept neighborhood and dual additions

philipc2 · 2025-03-07T22:23:17Z

HI @ahijevyc

Apologies for not getting to this PR earlier.

Looking at the implementation here, it looks great. It does however bring to light a possible need for us to consider a better, more streamlined, approach to handling these types of groupings and then applying some function on the result.

I mention this because of our Topological Aggregations. For this family of functions, we have distinct methods (i.e. topological_mean()), which looking back at doesn't seem like the preferred approach, especially if we plan to implement groupings like the neighborhood one and perhaps other spatial ones.

Very generally speaking, these functions essentially:

Group unstructured grid elements based on some condition/algorithm. Here we use the KD/BallTree to determine the candidate elements, while in the topological aggregations we use the connectivity information
Apply some function to the grouping (i.e. mean())
Store the results back on the unstructured grid element (node, edge, or face)

I wonder if this would be a good opportunity to extend the inherited .groupby() method from Xarray to support these spatial groupings.

I'm not sure of calling these approaches "kernels" is appropriate, but for the sake of this example, we could provide spatial kernels the user could pass into groupby() and then perform aggregations directly on the result. This feels much more in line with Xarray's design philosophy.

# radial neighborhood of r=0.25
uxds['t2m'].groupby(kernel=ux.BoundingCircle(r=0.25)).mean()

# 2 deg by 2 deg bounding box 
uxds['t2m'].groupby(kernel=ux.BoundingBox(dlon=2, dlat=2))

# group the nodes that surround each face and find the maximum
uxds['node_centered_var'].groupby(kernel=ux.FaceNode()).max()

# this is equivalent to the following in the current release
uxds['node_centered_var'].topological_max(destination='face')

I'll ping @aaronzedwick and @erogluorhan for their thoughts on this. I personally really like the design above and think that it aligns well with the overall design.

aaronzedwick · 2025-03-10T16:34:00Z

HI @ahijevyc

Apologies for not getting to this PR earlier.

Looking at the implementation here, it looks great. It does however bring to light a possible need for us to consider a better, more streamlined, approach to handling these types of groupings and then applying some function on the result.

I mention this because of our Topological Aggregations. For this family of functions, we have distinct methods (i.e. topological_mean()), which looking back at doesn't seem like the preferred approach, especially if we plan to implement groupings like the neighborhood one and perhaps other spatial ones.

Very generally speaking, these functions essentially:

Group unstructured grid elements based on some condition/algorithm. Here we use the KD/BallTree to determine the candidate elements, while in the topological aggregations we use the connectivity information

Apply some function to the grouping (i.e. mean())

Store the results back on the unstructured grid element (node, edge, or face)

I wonder if this would be a good opportunity to extend the inherited .groupby() method from Xarray to support these spatial groupings.

I'm not sure of calling these approaches "kernels" is appropriate, but for the sake of this example, we could provide spatial kernels the user could pass into groupby() and then perform aggregations directly on the result. This feels much more in line with Xarray's design philosophy.
# radial neighborhood of r=0.25
uxds['t2m'].groupby(kernel=ux.BoundingCircle(r=0.25)).mean()

# 2 deg by 2 deg bounding box 
uxds['t2m'].groupby(kernel=ux.BoundingBox(dlon=2, dlat=2))

# group the nodes that surround each face and find the maximum
uxds['node_centered_var'].groupby(kernel=ux.FaceNode()).max()

# this is equivalent to the following in the current release
uxds['node_centered_var'].topological_max(destination='face')
I'll ping @aaronzedwick and @erogluorhan for their thoughts on this. I personally really like the design above and think that it aligns well with the overall design.

That is interesting. You suggesting changing the way we do aggregations entirely? Then this would affect the reduction PR I am working on then. Perhaps this PR could implement that change if you wish. I am fine with this, if you want to, it sounds like it would be intuitive.

philipc2 · 2025-03-10T17:53:25Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.

The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

aaronzedwick · 2025-03-10T20:03:00Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.

The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

So would the reductions PR be obsolete?

philipc2 · 2025-03-10T20:05:53Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.
The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

So would the reductions PR be obsolete?

No. The underlying implementation would remain the same, since we would still need those implemented.

This would just provide a different interface for it, with a more "Xarray-like" interface.

aaronzedwick · 2025-03-10T20:15:35Z

You suggesting changing the way we do aggregations entirely?

We should be changing the API for it, since the current one is not sustainable if we plan to implement many times of spatial grouping methods.
The underlying functionality would remain mostly unchanged.

Perhaps this PR could implement that change if you wish.

If we do decide to do this, I'll open up a separate PR starting with the topological aggregations.

So would the reductions PR be obsolete?

No. The underlying implementation would remain the same, since we would still need those implemented.

This would just provide a different interface for it, with a more "Xarray-like" interface.

Ah, okay, I see. That makes sense, thanks for the clarification!

rajeeja · 2026-07-28T18:29:40Z

@ahijevyc — since GitHub blocks adding the PR author as a reviewer, could you take a look when you get a chance? Key changes since your last review round:

Bug fixes:

np.empty → np.full(…, np.nan) in neighbors.py — empty neighborhoods now correctly return NaN instead of arbitrary memory
Auto-transpose for non-last grid dim — (time, n_face) arrays now work without the caller having to manually transpose
raise ValueError instead of bare assert for dimension mismatch (cleaner user-facing error)
_copy(data=…, deep=False) in the filter path — skips a redundant deep copy of the source data, shares the same uxgrid reference

API compliance / xarray parity:

Investigated whether neighborhood_filter should use groupby: conclusion is no. Neighborhoods are overlapping spatial balls — a single element belongs to many neighborhoods. xarray's groupby requires non-overlapping partitions. The direct-method pattern is consistent with azimuthal_mean and zonal_mean.
Philip has been removed as a reviewer; @erogluorhan has been added.

Docs:

Expanded user-guide notebook (30 cells): covers face-, node-, and edge-centered data; multi-dim (time × space); xarray chaining (.where(), .groupby()); radius sweep comparison; ux.tutorial datasets throughout (no hardcoded file paths)
Added Examples and See Also sections to both UxDataArray.neighborhood_filter and UxDataset.neighborhood_filter docstrings

Happy to discuss further or adjust anything before merge!

Philip is no longer on the project. Dismissing stale review — all concerns have been addressed in subsequent commits.

ahijevyc

Thanks for carrying this across the finish line! I just had a comment or two about the assertion that NaN could be returned for a coarse grid and small radius. I don't think this could occur in practice, as the filter should at least return the values at the original grid points even if the radius is zero. And grid.neighbors.BallTree.query_radius already makes sure the radius is not negative.

…_filter

Neighborhoods are never empty: query_radius rejects a negative radius and every element is its own neighbor at distance 0, so r=0 returns the original values. Reword the code comment and user-guide section accordingly, keeping np.full(NaN) allocation only as a defensive measure over np.empty. Also raise DataCenteringError instead of a bare ValueError for non-grid-mapped data, matching the error types introduced in uxarray/errors.py, and drop a redundant self-import in UxDataArray.isel.

- Include distance_metric in the get_ball_tree/get_kd_tree cache invalidation check. Only coordinates and coordinate_system were compared, so requesting a different metric silently returned a tree built with the original one. - Request a spherical/haversine tree explicitly in _neighborhood_filter. It previously read coordinate_system back off whatever tree was cached, so a cartesian tree from an earlier call made r a chord length instead of the documented great-circle degrees. - Raise an explanatory TypeError when func does not accept an axis keyword, instead of surfacing a raw NumPy message. - Document radius units, overlapping neighborhoods, and eager evaluation of dask-backed input in both public docstrings. - Restore the original _copy deep default; changing it was unrelated to this feature and affects every UxDataArray copy. Adds coverage for tree cache invalidation, the spherical-tree guarantee, the func-without-axis error, and dask input.

ahijevyc

looks great

…_filter

erogluorhan

Overall this looks great, but I have a couple concerns in the low-level implementation for which @cmdupuis3 could be helpful as well (added as a reviewer). See below

erogluorhan · 2026-08-04T19:54:46Z

+        func: Callable = np.mean,
+        r: float = 1.0,
+    ) -> UxDataArray:
+        """Apply a neighborhood filter, replacing the value at each grid


Might use a little rephrasing: "replace" might be misleading here since function returns a new data array and actually doesn't change the original data?

erogluorhan · 2026-08-04T19:58:56Z

+
+    # Apply func along the last (grid) axis only, so any extra leading
+    # dimensions (e.g. time) are preserved rather than being collapsed.
+    for i, idx in enumerate(neighbor_indices):


This will iterate through every single geometric element, e.g. faces, right? It looks very costly for km-scale data. Was there any consideration of vectorization for this?

cc: @cmdupuis3 for your thoughts

Something you could do is make a boolean mask out of neighbor_indices as keys, then mask it to neighbor_masked, then you could delete the whole conditional/try/except block and iterate over only the valid entries. If you really need the full-scale destination_data, you can use the mask to reshape the output.

erogluorhan · 2026-08-04T20:08:03Z

+    for i, idx in enumerate(neighbor_indices):
+        if len(idx):
+            try:
+                destination_data[..., i] = func(data[..., idx], axis=-1)


From what I can tell on this line, every single iteration will trigger rebuilding and executing the dask task graph from scratch because this line returns a lazy dask scalar when data is dask-backed and assigning that into a numpy array index.

I believe this could be avoided if uxdataarray.py/neighborhood_filter did an upfront .compute() for data (and actually that function already claims "lazy (dask-backed) data is computed eagerly and the result is always NumPy-backed").

cc: @cmdupuis3 you dealt with bundled .compute()s through xarray API in #1560, it might not directly related to here, but you may still have some thoughts.

Looks like it. Just at a glance, this block seems like it's basically a ufunc implementation.

cmdupuis3 · 2026-08-05T17:19:16Z

I'm working on a chunk-friendly optimized version based on gufuncs, I'll add my own PR off this one and y'all can merge it. The one caveat is that we'd have to hand-enumerate the specific reductions we want to be able to run optimally, I think otherwise we can offer a fallback option but the performance would suffer. I don't think that would happen that often in practice though.

rajeeja · 2026-08-05T18:15:36Z

I'm working on a chunk-friendly optimized version based on gufuncs, I'll add my own PR off this one and y'all can merge it. The one caveat is that we'd have to hand-enumerate the specific reductions we want to be able to run optimally, I think otherwise we can offer a fallback option but the performance would suffer. I don't think that would happen that often in practice though.

Makes sense, thanks for finding this out, that code would lose out-of-core/parallel benefits; most of the times this would be one of the steps in the bigger pipeline and if others are nice dask ops that would be a good win. I'll keep an eye on the followup to this.

Follows the gufunc work with the API changes it made possible. Kernel factory. The buffered kernels (median and friends) share one gather loop, with the reduction supplied as a numba-compilable callable. numba supports the NumPy reductions in nopython mode, so np.median is used directly rather than reimplemented -- it partitions instead of fully sorting, making it 1.6x faster than the hand-written sort it replaces, as well as shorter. Adds ptp, std, var, quantile and percentile; the parameterized ones carry their argument as a scalar core dimension so q and ddof vary per call without recompiling. Named reductions. `func` now takes a name ("mean", "quantile", ...) with parameters as ordinary keyword arguments. The previous `axis=-1` contract could not be jitted -- numba rejects the axis kwarg on mean/median/percentile -- and dispatching on a function object cannot see through functools.partial, so `partial(np.percentile, q=90)`, the example in our own docstring, could never reach a kernel. A name always can. Callables remain accepted on the old contract as an escape hatch, and np.mean and friends still map to their kernels so existing code keeps the fast path. Neighborhoods. Finding the neighbors costs more than reducing over them -- after the kernel work it is ~95% of a call -- and it was repeated on every call, including once per variable in the dataset path. Grid.neighborhoods() does the search once and returns an object to reduce over repeatedly: 3.5x for four reductions at one radius on a 196k-face grid, and a dataset filter now costs one query per grid location rather than one per variable (8 variables: 1.74s -> 0.23s). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

cmdupuis3 · 2026-08-05T23:57:47Z

@ahijevyc @rajeeja I have a draft PR sitting in the fork, I'll keep working on it there. I'm not totally satisfied with the level of complexity. Using a functional approach inside vectorized gufuncs makes it hard to get performance and good readability at once.

- numba guvectorize kernels for compiled parallel reductions - Neighborhoods object: one BallTree query reused across variables/reductions - Named reduction API: func='mean', func='percentile', q=90, etc.

cmdupuis3 · 2026-08-06T15:18:05Z

pre-commit.ci autofix

for more information, see https://pre-commit.ci

ahijevyc added 2 commits September 7, 2024 11:09

Neighborhood filter

75d832e

Merge remote-tracking branch 'upstream/main' into neighborhood_filter

b850bc4

ahijevyc added the new feature New feature or request label Sep 9, 2024

ahijevyc self-assigned this Sep 9, 2024

ahijevyc mentioned this pull request Sep 9, 2024

Add neighborhood filter to remap #923

Closed

14 tasks

ahijevyc added 2 commits September 9, 2024 10:48

ruff recommendations

8ec0193

added Callable to Type checking

5605949

philipc2 linked an issue Sep 9, 2024 that may be closed by this pull request

Apply a neighborhood filter with radius r to all elements of UxDataArray #930

Open

philipc2 reviewed Sep 9, 2024

View reviewed changes

Comment thread uxarray/core/dataarray.py Outdated

Merge branch 'main' into ahijevyc/neighborhood_filter

70ba961

philipc2 previously requested changes Sep 9, 2024

View reviewed changes

Comment thread uxarray/core/dataarray.py Outdated

Comment thread uxarray/core/dataarray.py

Comment thread uxarray/core/dataset.py Outdated

ahijevyc and others added 4 commits September 9, 2024 15:10

np.vstack().T faster than np.c

0c7bc1e

Fix some comments

d6d8a33

Merge branch 'main' into ahijevyc/neighborhood_filter

47b9cda

Merge branch 'main' into ahijevyc/neighborhood_filter

f75db0d

ahijevyc commented Oct 2, 2024

View reviewed changes

Comment thread uxarray/core/dataarray.py Outdated

Merge branch 'main' into ahijevyc/neighborhood_filter

bddd2fa

Kept neighborhood and dual additions

ahijevyc and others added 4 commits March 17, 2025 14:12

Merge branch 'main' into ahijevyc/neighborhood_filter

a0b6361

missing imports

6c59af7

Merge branch 'main' into ahijevyc/neighborhood_filter

67d0c11

Merge branch 'main' into ahijevyc/neighborhood_filter

8979745

philipc2 mentioned this pull request May 14, 2025

Add Azimuthal Averaging #1255

Merged

9 tasks

ahijevyc added 2 commits May 28, 2025 11:47

Merge branch 'main' into ahijevyc/neighborhood_filter

a8875cd

Merge branch 'UXARRAY:main' into ahijevyc/neighborhood_filter

f4af498

rajeeja added 2 commits July 28, 2026 13:54

Merge branch 'main' into ahijevyc/neighborhood_filter

29f1873

Fix pre-commit: trailing newlines + ruff import sort/format

b37974d

Merge branch 'UXARRAY:main' into ahijevyc/neighborhood_filter

ed35c01

ahijevyc commented Jul 29, 2026

View reviewed changes

Comment thread docs/user-guide/neighborhood-filter.ipynb

ahijevyc commented Jul 29, 2026

View reviewed changes

Comment thread uxarray/grid/neighbors.py Outdated

rajeeja added 3 commits August 3, 2026 08:46

Merge remote-tracking branch 'origin/main' into ahijevyc/neighborhood…

a2aff9c

…_filter

ahijevyc commented Aug 3, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into ahijevyc/neighborhood…

5f33a72

…_filter

erogluorhan requested review from philipc2 and removed request for philipc2 August 4, 2026 19:37

erogluorhan reviewed Aug 4, 2026

View reviewed changes

erogluorhan requested a review from cmdupuis3 August 4, 2026 20:19

cmdupuis3 reviewed Aug 5, 2026

View reviewed changes

Comment thread uxarray/core/dataarray.py Outdated

cmdupuis3 and others added 3 commits August 5, 2026 13:40

941: gufunc for neighborhood filter

b964260

941: some cleanup

c7d2711

rajeeja added 2 commits August 6, 2026 02:32

Merge main into ahijevyc/neighborhood_filter to resolve conflicts

c258e63

Merge Chris Dupuis gufunc/Neighborhoods work into neighborhood_filter

39479a1

- numba guvectorize kernels for compiled parallel reductions - Neighborhoods object: one BallTree query reused across variables/reductions - Named reduction API: func='mean', func='percentile', q=90, etc.

rajeeja assigned cmdupuis3 Aug 6, 2026

pre-commit-ci Bot and others added 2 commits August 6, 2026 15:18

[pre-commit.ci] auto fixes from pre-commit.com hooks

c240cd3

for more information, see https://pre-commit.ci

Merge branch 'main' into ahijevyc/neighborhood_filter

728bf87

Uh oh!

Conversation

ahijevyc commented Sep 9, 2024 • edited by rajeeja Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Expected Usage

PR Checklist

Uh oh!

Uh oh!

philipc2 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

philipc2 commented Mar 7, 2025

Uh oh!

aaronzedwick commented Mar 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

philipc2 commented Mar 10, 2025

Uh oh!

aaronzedwick commented Mar 10, 2025

Uh oh!

philipc2 commented Mar 10, 2025

Uh oh!

aaronzedwick commented Mar 10, 2025

Uh oh!

rajeeja commented Jul 28, 2026

Uh oh!

Uh oh!

ahijevyc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ahijevyc left a comment

Choose a reason for hiding this comment

Uh oh!

erogluorhan left a comment

Choose a reason for hiding this comment

Uh oh!

erogluorhan Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

erogluorhan Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

cmdupuis3 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

erogluorhan Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

cmdupuis3 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cmdupuis3 commented Aug 5, 2026

Uh oh!

rajeeja commented Aug 5, 2026

Uh oh!

cmdupuis3 commented Aug 5, 2026

Uh oh!

cmdupuis3 commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ahijevyc commented Sep 9, 2024 •

edited by rajeeja

Loading

aaronzedwick commented Mar 10, 2025 •

edited

Loading